import sys
import time

import torch
from eagle.model.ea_model import EaModel
from fastchat.model import get_conversation_template
base_model_path = '/model/Qwen2-0.5B-Instruct'
EAGLE_model_path = '/model/eagle-0.5B-head-14'
# EAGLE_model_path = '/model/eagle-0.5B-head-14'
model = EaModel.from_pretrained(
    base_model_path=base_model_path,
    ea_model_path=EAGLE_model_path,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    device_map="auto",
    total_token=-1
)
model.eval()
questions=['quit',
'你是',
'我很高兴',]

# warnmup
for i, your_message in enumerate(questions):  # 遍历问题列表
    start = time.perf_counter()
    try:
        conv = get_conversation_template("qwen")
        conv.append_message(conv.roles[0], your_message)
        conv.append_message(conv.roles[1], None)
        prompt = conv.get_prompt()
        input_ids = model.tokenizer([prompt]).input_ids
        if torch.cuda.is_available():  # 检查 CUDA 是否可用
            input_ids = torch.as_tensor(input_ids).cuda()
        else:
            input_ids = torch.as_tensor(input_ids)
        output_ids = model.eagenerate(input_ids, temperature=0.5, max_new_tokens=512)
        output = model.tokenizer.decode(output_ids[0])
        end = time.perf_counter()
        print('')
    except Exception as e:
        break

# questions=['月亮是什么',
# '嫦娥在月亮上吗',
# '讲两个不同的笑话用序号表示',
# '模型预热',
# '你是谁',
# '你预热完成了没有',
# '请告诉我你最重要的技能',
# '写个 python 类实现斐波那契数列算法',
# '写一首诗要求五言绝句',
# '蛇年农历春节马上要到了，请帮我写一副很屌的对联']

questions=['某公司在未进行数据出境安全评估的情况下，将大量包含个人敏感信息的数据直接出口到境外，而且并未与境外接收方签订关于数据安全保护的法律文件。请根据以下法规简要回答这一行为是否违法。《数据出境安全评估办法》 涉及风险领域：数据治理 涉及研发条款： 第四条 数据处理者向境外提供数据，有下列情形之一的，应当通过所在地省级网信部门向国家网信部门申报数据出境安全评估： （一）数据处理者向境外提供重要数据； （二）关键信息基础设施运营者和处理100万人以上个人信息的数据处理者向境外提供个人信息； （三）自上年1月1日起累计向境外提供10万人个人信息或者1万人敏感个人信息的数据处理者向境外提供个人信息； （四）国家网信部门规定的其他需要申报数据出境安全评估的情形。 第五条 数据处理者在申报数据出境安全评估前，应当开展数据出境风险自评估，重点评估以下事项： （一）数据出境和境外接收方处理数据的目的、范围、方式等的合法性、正当性、必要性； （二）出境数据的规模、范围、种类、敏感程度，数据出境可能对国家安全、公共利益、个人或者组织合法权益带来的风险； （三）境外接收方承诺承担的责任义务，以及履行责任义务的管理和技术措施、能力等能否保障出境数据的安全； （四）数据出境中和出境后遭到篡改、破坏、泄露、丢失、转移或者被非法获取、非法利用等的风险，个人信息权益维护的渠道是否通畅等； （五）与境外接收方拟订立的数据出境相关合同或者其他具有法律效力的文件等（以下统称法律文件）是否充分约定了数据安全保护责任义务； （六）其他可能影响数据出境安全的事项。 第九条 数据处理者应当在与境外接收方订立的法律文件中明确约定数据安全保护责任义务，至少包括以下内容： （一）数据出境的目的、方式和数据范围，境外接收方处理数据的用途、方式等； （二）数据在境外保存地点、期限，以及达到保存期限、完成约定目的或者法律文件终止后出境数据的处理措施； （三）对于境外接收方将出境数据再转移给其他组织、个人的约束性要求； （四）境外接收方在实际控制权或者经营范围发生实质性变化，或者所在国家、地区数据安全保护政策法规和网络安全环境发生变化以及发生其他不可抗力情形导致难以保障数据安全时，应当采取的安全措施； （五）违反法律文件约定的数据安全保护义务的补救措施、违约责任和争议解决方式； （六）出境数据遭到篡改、破坏、泄露、丢失、转移或者被非法获取、非法利用等风险时，妥善开展应急处置的要求和保障个人维护其个人信息权益的途径和方式。 第十四条 通过数据出境安全评估的结果有效期为2年，自评估结果出具之日起计算。在有效期内出现以下情形之一的，数据处理者应当重新申报评估： （一）向境外提供数据的目的、方式、范围、种类和境外接收方处理数据的用途、方式发生变化影响出境数据安全的，或者延长个人信息和重要数据境外保存期限的； （二）境外接收方所在国家或者地区数据安全保护政策法规和网络安全环境发生变化以及发生其他不可抗力情形、数据处理者或者境外接收方实际控制权发生变化、数据处理者与境外接收方法律文件变更等影响出境数据安全的； （三）出现影响出境数据安全的其他情形。 有效期届满，需要继续开展数据出境活动的，数据处理者应当在有效期届满60个工作日前重新申报评估。 第十九条 本办法所称重要数据，是指一旦遭到篡改、破坏、泄露或者非法获取、非法利用等，可能危害国家安全、经济运行、社会稳定、公共健康和安全等的数据',
           '某公司将大量个人敏感信息发送至境外接收方，在与境外接收方订立的法律文件中并未明确约定数据安全保护责任义务，包括数据出境的目的、方式和数据范围，境外接收方处理数据的用途、方式等，以及境外保存地点、期限，以及达到保存期限、完成约定目的或者法律文件终止后出境数据的处理措施等内容。请根据以下法规简要回答这一行为是否违法。《数据出境安全评估办法》 涉及风险领域：数据治理 涉及研发条款： 第四条 数据处理者向境外提供数据，有下列情形之一的，应当通过所在地省级网信部门向国家网信部门申报数据出境安全评估： （一）数据处理者向境外提供重要数据； （二）关键信息基础设施运营者和处理100万人以上个人信息的数据处理者向境外提供个人信息； （三）自上年1月1日起累计向境外提供10万人个人信息或者1万人敏感个人信息的数据处理者向境外提供个人信息； （四）国家网信部门规定的其他需要申报数据出境安全评估的情形。 第五条 数据处理者在申报数据出境安全评估前，应当开展数据出境风险自评估，重点评估以下事项： （一）数据出境和境外接收方处理数据的目的、范围、方式等的合法性、正当性、必要性； （二）出境数据的规模、范围、种类、敏感程度，数据出境可能对国家安全、公共利益、个人或者组织合法权益带来的风险； （三）境外接收方承诺承担的责任义务，以及履行责任义务的管理和技术措施、能力等能否保障出境数据的安全； （四）数据出境中和出境后遭到篡改、破坏、泄露、丢失、转移或者被非法获取、非法利用等的风险，个人信息权益维护的渠道是否通畅等； （五）与境外接收方拟订立的数据出境相关合同或者其他具有法律效力的文件等（以下统称法律文件）是否充分约定了数据安全保护责任义务； （六）其他可能影响数据出境安全的事项。 第九条 数据处理者应当在与境外接收方订立的法律文件中明确约定数据安全保护责任义务，至少包括以下内容： （一）数据出境的目的、方式和数据范围，境外接收方处理数据的用途、方式等； （二）数据在境外保存地点、期限，以及达到保存期限、完成约定目的或者法律文件终止后出境数据的处理措施； （三）对于境外接收方将出境数据再转移给其他组织、个人的约束性要求； （四）境外接收方在实际控制权或者经营范围发生实质性变化，或者所在国家、地区数据安全保护政策法规和网络安全环境发生变化以及发生其他不可抗力情形导致难以保障数据安全时，应当采取的安全措施； （五）违反法律文件约定的数据安全保护义务的补救措施、违约责任和争议解决方式； （六）出境数据遭到篡改、破坏、泄露、丢失、转移或者被非法获取、非法利用等风险时，妥善开展应急处置的要求和保障个人维护其个人信息权益的途径和方式。 第十四条 通过数据出境安全评估的结果有效期为2年，自评估结果出具之日起计算。在有效期内出现以下情形之一的，数据处理者应当重新申报评估： （一）向境外提供数据的目的、方式、范围、种类和境外接收方处理数据的用途、方式发生变化影响出境数据安全的，或者延长个人信息和重要数据境外保存期限的； （二）境外接收方所在国家或者地区数据安全保护政策法规和网络安全环境发生变化以及发生其他不可抗力情形、数据处理者或者境外接收方实际控制权发生变化、数据处理者与境外接收方法律文件变更等影响出境数据安全的； （三）出现影响出境数据安全的其他情形。 有效期届满，需要继续开展数据出境活动的，数据处理者应当在有效期届满60个工作日前重新申报评估。 第十九条 本办法所称重要数据，是指一旦遭到篡改、破坏、泄露或者非法获取、非法利用等，可能危害国家安全、经济运行、社会稳定、公共健康和安全等的数据。", "input": "", "output": "根据《数据出境安全评估办法》，某公司将大量个人敏感信息发送至境外接收方且未在法律文件中明确约定数据安全保护责任义务是违法的。根据该办法，数据处理者应当在与境外接收方订立的法律文件中明确约定数据安全保护责任义务，包括数据出境的目的、方式和数据范围，境外接收方处理数据的用途、方式等，以及境外保存地点、期限，以及达到保存期限、完成约定目的或者法律文件终止后出境数据的处理措施等内容。违反这些规定将会导致违法行为。']


for i, your_message in enumerate(questions):  # 遍历问题列表
    start = time.perf_counter()
    try:
        conv = get_conversation_template("qwen")
        conv.append_message(conv.roles[0], your_message)
        conv.append_message(conv.roles[1], None)
        prompt = conv.get_prompt()
        input_ids = model.tokenizer([prompt]).input_ids
        if torch.cuda.is_available():  # 检查 CUDA 是否可用
            input_ids = torch.as_tensor(input_ids).cuda()
        else:
            input_ids = torch.as_tensor(input_ids)
        output_ids = model.eagenerate(input_ids, temperature=0.5, max_new_tokens=512)
        output = model.tokenizer.decode(output_ids[0])
        end = time.perf_counter()
        print(f'Question {i + 1}: {your_message}总耗时{end - start},总token{len(output)},单个token平均耗时{(end - start) / len(output)}')
        print(f'Answer:{output}')
        print('')
    except Exception as e:
        print(f"An error occurred: {e}")
        break
